# build_requestion.py
# usage: 根据已经标记0和-1的check文件，生成requestion文件
# python ./src/tools/build_requestion.py --path "./data" --requestion "./data/requestion.jsonl"

import os
import json
import argparse
from pathlib import Path

def main():
    parser = argparse.ArgumentParser(description="扫描指定路径下以 'check' 开头的文件，提取 flag=0 或 flag=-1 的记录，写入 requestion 文件。")
    parser.add_argument('--path', type=str, required=True, help='包含 check 开头文件的目录路径')
    parser.add_argument('--requestion', type=str, required=True, help='输出文件路径（含文件名）')

    args = parser.parse_args()

    check_path = Path(args.path)
    output_file = Path(args.requestion)

    if not check_path.exists():
        print(f"错误：路径 '{check_path}' 不存在。")
        return

    if not check_path.is_dir():
        print(f"错误：'{check_path}' 不是一个有效的目录。")
        return

    # 获取所有以 'check' 开头的文件
    check_files = [f for f in check_path.iterdir() if f.is_file() and f.name.startswith('check')]

    if not check_files:
        print(f"警告：在 '{check_path}' 中未找到以 'check' 开头的文件。")
        return

    output_records = []

    for file_path in check_files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        data = json.loads(line)
                        # 检查是否包含 flag 字段且为 0 或 -1
                        if 'flag' in data and data['flag'] in (0, -1):
                            # 提取需要的字段
                            record = {
                                "id": data.get("id"),
                                "question": data.get("question"),
                                "sql": data.get("sql"),
                                "level": data.get("level")
                            }
                            output_records.append(record)
                    except json.JSONDecodeError as e:
                        print(f"文件 {file_path.name} 第 {line_num} 行 JSON 解析错误: {e}")
        except Exception as e:
            print(f"读取文件 {file_path.name} 时发生错误: {e}")

    # # 写入输出文件
    # try:
    #     # 确保输出文件的目录存在
    #     output_file.parent.mkdir(parents=True, exist_ok=True)

    #     with open(output_file, 'w', encoding='utf-8') as f:
    #         for record in output_records:
    #             f.write(json.dumps(record, ensure_ascii=False) + '\n')
    #     print(f"成功：共提取 {len(output_records)} 条记录，已写入 '{output_file}'。")
    # except Exception as e:
    #     print(f"写入文件 '{output_file}' 时发生错误: {e}")
        # 写入输出文件
    try:
        # 确保输出文件的目录存在
        output_file.parent.mkdir(parents=True, exist_ok=True)

        # === 新增：按 id 字段排序（支持字符串和数字 ID）===
        # 如果 id 是数字（int），可以转换为 int 排序
        # 如果是字符串，则直接字符串排序
        # 这里使用健壮的方式：尝试转为 int，失败则用字符串排序
        def get_sort_key(record):
            _id = record.get("id")
            if not _id:
                return ("", "")  # 处理缺失 id 的情况
            try:
                return (0, int(_id))  # 数字优先
            except ValueError:
                return (1, str(_id))  # 非数字次之

        output_records.sort(key=get_sort_key)

        # ===================================================

        with open(output_file, 'w', encoding='utf-8') as f:
            for record in output_records:
                f.write(json.dumps(record, ensure_ascii=False) + '\n')
        print(f"成功：共提取 {len(output_records)} 条记录，已按 id 排序后写入 '{output_file}'。")
    except Exception as e:
        print(f"写入文件 '{output_file}' 时发生错误: {e}")

if __name__ == '__main__':
    main()